package com.zenithmind.news.websocket;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketSession;

import java.time.LocalDateTime;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * WebSocket会话管理器 - 遵循单一职责原则
 * 专门负责WebSocket会话的管理
 * 
 * @author ZenithMind Team
 * @since 2025-01-09
 */
@Slf4j
@Component
public class WebSocketSessionManager {

    // 存储所有WebSocket连接 - 用户ID -> WebSocket会话集合
    private final Map<String, CopyOnWriteArraySet<WebSocketSession>> userSessions = new ConcurrentHashMap<>();
    
    // 存储会话信息 - 会话ID -> 用户信息
    private final Map<String, UserSessionInfo> sessionInfo = new ConcurrentHashMap<>();

    /**
     * 添加用户会话
     */
    public void addUserSession(String userId, WebSocketSession session, String clientIp) {
        userSessions.computeIfAbsent(userId, k -> new CopyOnWriteArraySet<>()).add(session);
        
        UserSessionInfo info = new UserSessionInfo();
        info.setUserId(userId);
        info.setClientIp(clientIp);
        info.setConnectTime(LocalDateTime.now());
        sessionInfo.put(session.getId(), info);
        
        log.info("添加用户会话：用户ID={}, 会话ID={}, IP={}", userId, session.getId(), clientIp);
    }

    /**
     * 移除用户会话
     */
    public void removeUserSession(WebSocketSession session) {
        UserSessionInfo info = sessionInfo.remove(session.getId());
        if (info != null) {
            String userId = info.getUserId();
            CopyOnWriteArraySet<WebSocketSession> sessions = userSessions.get(userId);
            if (sessions != null) {
                sessions.remove(session);
                if (sessions.isEmpty()) {
                    userSessions.remove(userId);
                }
            }
            log.info("移除用户会话：用户ID={}, 会话ID={}", userId, session.getId());
        }
    }

    /**
     * 获取用户的所有会话
     */
    public CopyOnWriteArraySet<WebSocketSession> getUserSessions(String userId) {
        return userSessions.get(userId);
    }

    /**
     * 获取会话信息
     */
    public UserSessionInfo getSessionInfo(String sessionId) {
        return sessionInfo.get(sessionId);
    }

    /**
     * 更新心跳时间
     */
    public void updateHeartbeat(String sessionId) {
        UserSessionInfo info = sessionInfo.get(sessionId);
        if (info != null) {
            info.setLastHeartbeatTime(LocalDateTime.now());
        }
    }

    /**
     * 获取所有用户会话
     */
    public Map<String, CopyOnWriteArraySet<WebSocketSession>> getAllUserSessions() {
        return userSessions;
    }

    /**
     * 用户会话信息
     */
    public static class UserSessionInfo {
        private String userId;
        private String clientIp;
        private LocalDateTime connectTime;
        private LocalDateTime lastHeartbeatTime;

        // Getters and Setters
        public String getUserId() { return userId; }
        public void setUserId(String userId) { this.userId = userId; }
        
        public String getClientIp() { return clientIp; }
        public void setClientIp(String clientIp) { this.clientIp = clientIp; }
        
        public LocalDateTime getConnectTime() { return connectTime; }
        public void setConnectTime(LocalDateTime connectTime) { this.connectTime = connectTime; }
        
        public LocalDateTime getLastHeartbeatTime() { return lastHeartbeatTime; }
        public void setLastHeartbeatTime(LocalDateTime lastHeartbeatTime) { this.lastHeartbeatTime = lastHeartbeatTime; }
    }
}
